package com.yupi.oj.service.impl;


import cn.hutool.core.collection.CollUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.yupi.oj.common.ErrorCode;
import com.yupi.oj.common.InvokeAiNumber;
import com.yupi.oj.exception.BusinessException;
import com.yupi.oj.exception.ThrowUtils;
import com.yupi.oj.manager.AiManager;
import com.yupi.oj.mapper.AiAnswerMapper;
import com.yupi.oj.model.dto.ai.AiMessage;
import com.yupi.oj.model.dto.ai.AiRequest;
import com.yupi.oj.model.entity.AiAnswer;
import com.yupi.oj.model.entity.Question;
import com.yupi.oj.model.entity.RecordAiNumber;
import com.yupi.oj.model.entity.User;
import com.yupi.oj.model.enums.AiRoleEnum;
import com.yupi.oj.model.enums.UserRoleEnum;
import com.yupi.oj.service.AiAnswerService;
import com.yupi.oj.service.QuestionService;
import com.yupi.oj.service.RecordAiNumberService;
import com.yupi.oj.service.UserService;
import io.reactivex.Flowable;
import io.reactivex.schedulers.Schedulers;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletRequest;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * @author 兰飞宇
 * @description 针对表【ai_answer(AI回答表)】的数据库操作Service实现
 * @createDate 2025-03-20 16:40:06
 */
@Service
@Slf4j
public class AiAnswerServiceImpl extends ServiceImpl<AiAnswerMapper, AiAnswer>
        implements AiAnswerService {
    @Autowired
    private QuestionService questionService;
    @Autowired
    private AiManager aiManager;
    @Autowired
    private UserService userService;
    @Autowired
    private RecordAiNumberService recordAiNumberService;
    @Autowired
    private RedisTemplate redisTemplate;
    private String SYSTEM_MESSAGE;
    private static final String REDIS_KEY = "oj:ai:questionId:userId";


    public String getSystemMessage(String title, String content) {
        return String.format("你是一名编程专家，你需要解答的题目如下：\n" +
                "```\n" +
                "题目标题：%s\n" +
                "题目内容：【【【%s】】】\n" +
                "```\n" +
                "请根据以上题目给用户解答问题。\n" +
                "要求：\n" +
                "1.用户如果问了跟这道题目无关的问题，那么统一回复“抱歉，我仅为您解答这道题目，请不要问无关的问题”\n" +
                "2.给出的解答尽量精简", title, content);
    }

    public void fillSystemMessage(Long questionId) {
        Question question = questionService.getById(questionId);
        ThrowUtils.throwIf(question == null, ErrorCode.NOT_FOUND_ERROR, "题目不存在");
        String title = question.getTitle();
        String content = question.getContent();
        //去掉content中的多余空行和特殊字符
        content = content.replaceAll("(\n\\s*)+", "\n");
        content = content.replace("#", "");
        SYSTEM_MESSAGE = getSystemMessage(title, content);
    }

    /**
     * 封装redisKey
     * @return
     */
    public String getRedisKey(Long questionId,Long userId){
        return String.format("%s:%s:%s", REDIS_KEY, questionId, userId);
    }

    /**
     * 非流式返回结果
     * @param aiRequest
     * @param request
     * @return
     */
    @Override
    public String doChat(AiRequest aiRequest, HttpServletRequest request) {
        //校验是否登录
        User loginUser = userService.getLoginUser(request);
        ThrowUtils.throwIf(loginUser == null, ErrorCode.NOT_LOGIN_ERROR);
        Long userId = loginUser.getId();
        Long questionId = aiRequest.getQuestionId();
        String userMessage = aiRequest.getMessage();
        //查看用户调用AI记录表，如果调用次数超过10，那么直接返回
        LambdaQueryWrapper<RecordAiNumber> queryWrapper = new LambdaQueryWrapper<RecordAiNumber>()
                .eq(RecordAiNumber::getUserId, userId);
        RecordAiNumber serviceOne = recordAiNumberService.getOne(queryWrapper);
        //如果数据存在并且权限是普通用户
        if (serviceOne != null && serviceOne.getUserRole().equals(UserRoleEnum.USER.getValue())) {
            Integer number = serviceOne.getNumber();
            ThrowUtils.throwIf(number >= InvokeAiNumber.DEFAULT_USER_NUMBER, ErrorCode.OPERATION_ERROR, "您今日的AI调用次数已经用完，请明天再来");
        }
        //查询ai_answer表
        LambdaQueryWrapper<AiAnswer> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        LambdaQueryWrapper<AiAnswer> eq = lambdaQueryWrapper
                .eq(AiAnswer::getQuestionId, questionId)
                .eq(AiAnswer::getUserId, userId);
        AiAnswer aiAnswer = this.getOne(eq);
        //如果aiAnswer表存在，封装数据, 不存在填充prompt
        List<AiMessage> aiMessageList = new ArrayList<>();
        if (aiAnswer != null) {
            String message = aiAnswer.getMessage();
            aiMessageList = JSONUtil.toList(message, AiMessage.class);
        } else {
            //填充prompt
            fillSystemMessage(questionId);
            //将systemMessage加入到集合中
            aiMessageList.add(new AiMessage(AiRoleEnum.SYSTEM.getValue(), SYSTEM_MESSAGE));
            log.info("SYSTEM_MESSAGE:{}", SYSTEM_MESSAGE);
        }
        List<String> userMessageList = new ArrayList<>();
        List<String> assistantMessageList = new ArrayList<>();
        List<String> systemMessageList = new ArrayList<>();
        //如果aiMessageList不为空且表中数据存在，封装传给AI服务的参数
        if (aiAnswer != null && CollUtil.isNotEmpty(aiMessageList)) {
            for (AiMessage aiMessage : aiMessageList) {
                if (aiMessage.getRole().equals(AiRoleEnum.USER.getValue())) {
                    userMessageList.add(aiMessage.getMessage());
                } else if (aiMessage.getRole().equals(AiRoleEnum.ASSISTANT.getValue())) {
                    assistantMessageList.add(aiMessage.getMessage());
                } else {
                    systemMessageList.add(aiMessage.getMessage());
                }
            }
        }
        userMessageList.add(userMessage);
        //将用户问题封装到aiMessageList中
        aiMessageList.add(new AiMessage(AiRoleEnum.USER.getValue(), userMessage));
        String result = "";
        //调用AI服务
        try {
            //如果数据还不存在，说明是第一次提问
            if (aiAnswer == null) {
                result = aiManager.doUnStreamAi(SYSTEM_MESSAGE, userMessage, null);
            } else {
                result = aiManager.doAnyUnStreamAi(systemMessageList.get(0), userMessageList, assistantMessageList, null);
            }
            aiMessageList.add(new AiMessage(AiRoleEnum.ASSISTANT.getValue(), result));
        } catch (Exception e) {
            e.printStackTrace();
            ThrowUtils.throwIf(true, ErrorCode.SYSTEM_ERROR, "AI服务异常");
        }
        //如果数据不存在，那么新增数据
        if (aiAnswer == null) {
            aiAnswer = new AiAnswer();
            aiAnswer.setQuestionId(questionId);
            aiAnswer.setUserId(userId);
            aiAnswer.setMessage(JSONUtil.toJsonStr(aiMessageList));
            boolean save = this.save(aiAnswer);
            ThrowUtils.throwIf(!save, ErrorCode.OPERATION_ERROR, "保存数据失败");
            //添加用户调用AI次数记录
            RecordAiNumber recordAiNumber = new RecordAiNumber();
            recordAiNumber.setUserId(userId);
            recordAiNumber.setUserRole(loginUser.getUserRole());
            recordAiNumber.setNumber(1);
            boolean save1 = recordAiNumberService.save(recordAiNumber);
            ThrowUtils.throwIf(!save1, ErrorCode.OPERATION_ERROR, "保存数据失败");
        } else {
            //数据存在更新数据
            aiAnswer.setMessage(JSONUtil.toJsonStr(aiMessageList));
            boolean b = this.updateById(aiAnswer);
            ThrowUtils.throwIf(!b, ErrorCode.OPERATION_ERROR, "更新数据失败");
            //更新用户调用AI次数记录
            if (serviceOne != null) {
                serviceOne.setNumber(serviceOne.getNumber() + 1);
                boolean updateById = recordAiNumberService.updateById(serviceOne);
                ThrowUtils.throwIf(!updateById, ErrorCode.OPERATION_ERROR, "更新数据失败");
            }
        }
        return result;
    }

    /**
     * 流式返回结果
     * @param aiRequest
     * @param request
     * @return
     */
    @Override
    public SseEmitter doChatStream(AiRequest aiRequest, HttpServletRequest request) {
        //校验是否登录
        User loginUser = userService.getLoginUser(request);
        ThrowUtils.throwIf(loginUser == null, ErrorCode.NOT_LOGIN_ERROR);
        Long userId = loginUser.getId();
        Long questionId = aiRequest.getQuestionId();
        String userMessage = aiRequest.getMessage();
        //查看用户调用AI记录表，如果调用次数超过10，那么直接返回
        LambdaQueryWrapper<RecordAiNumber> queryWrapper = new LambdaQueryWrapper<RecordAiNumber>()
                .eq(RecordAiNumber::getUserId, userId);
        RecordAiNumber serviceOne = recordAiNumberService.getOne(queryWrapper);
        //如果数据存在并且权限是普通用户
        if (serviceOne != null && serviceOne.getUserRole().equals(UserRoleEnum.USER.getValue())) {
            Integer number = serviceOne.getNumber();
            ThrowUtils.throwIf(number >= InvokeAiNumber.DEFAULT_USER_NUMBER, ErrorCode.OPERATION_ERROR, "您今日的AI调用次数已经用完，请明天再来");
        }
        //查询ai_answer表
        LambdaQueryWrapper<AiAnswer> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        LambdaQueryWrapper<AiAnswer> eq = lambdaQueryWrapper
                .eq(AiAnswer::getQuestionId, questionId)
                .eq(AiAnswer::getUserId, userId);
        final AiAnswer[] aiAnswer = {this.getOne(eq)};
        //如果aiAnswer表存在，封装数据, 不存在填充prompt
        List<AiMessage> aiMessageList;
        if (aiAnswer[0] != null) {
            String message = aiAnswer[0].getMessage();
            aiMessageList = JSONUtil.toList(message, AiMessage.class);
        } else {
            aiMessageList = new ArrayList<>();
            //填充prompt
            fillSystemMessage(questionId);
            //将systemMessage加入到集合中
            aiMessageList.add(new AiMessage(AiRoleEnum.SYSTEM.getValue(), SYSTEM_MESSAGE));
            log.info("SYSTEM_MESSAGE:{}", SYSTEM_MESSAGE);
        }
        List<String> userMessageList = new ArrayList<>();
        List<String> assistantMessageList = new ArrayList<>();
        List<String> systemMessageList = new ArrayList<>();
        //如果表中数据存在，封装传给AI服务的参数
        if (aiAnswer[0] != null) {
            for (AiMessage aiMessage : aiMessageList) {
                if (aiMessage.getRole().equals(AiRoleEnum.USER.getValue())) {
                    userMessageList.add(aiMessage.getMessage());
                } else if (aiMessage.getRole().equals(AiRoleEnum.ASSISTANT.getValue())) {
                    assistantMessageList.add(aiMessage.getMessage());
                } else {
                    systemMessageList.add(aiMessage.getMessage());
                }
            }
        }
        //将当前用户输入添加到集合中
        userMessageList.add(userMessage);
        //将用户问题封装到aiMessageList中
        aiMessageList.add(new AiMessage(AiRoleEnum.USER.getValue(), userMessage));
        SseEmitter sseEmitter = new SseEmitter(0L);
        StringBuilder stringBuilder = new StringBuilder();
        //调用AI服务
        try {
            Flowable<ChatCompletionChunk> flowable = aiManager.doAnyStreamAi(aiMessageList.get(0).getMessage(), userMessageList, assistantMessageList, null);
            flowable.observeOn(Schedulers.io())
                    .doOnNext(message -> {
                        Object content = message.getChoices().get(0).getMessage().getContent();
                        stringBuilder.append(content);
                        sseEmitter.send(content);
                    })
                    .doOnComplete(() -> {
                        aiMessageList.add(new AiMessage(AiRoleEnum.ASSISTANT.getValue(), stringBuilder.toString()));
                        //如果数据不存在，那么新增数据
                        if (aiAnswer[0] == null) {
                            aiAnswer[0] = new AiAnswer();
                            aiAnswer[0].setQuestionId(questionId);
                            aiAnswer[0].setUserId(userId);
                            aiAnswer[0].setMessage(JSONUtil.toJsonStr(aiMessageList));
                            boolean save = this.save(aiAnswer[0]);
                            ThrowUtils.throwIf(!save, ErrorCode.OPERATION_ERROR, "保存数据失败");
                            //添加用户调用AI次数记录
                            if(serviceOne == null){
                                RecordAiNumber recordAiNumber = new RecordAiNumber();
                                recordAiNumber.setUserId(userId);
                                recordAiNumber.setUserRole(loginUser.getUserRole());
                                recordAiNumber.setNumber(1);
                                boolean save1 = recordAiNumberService.save(recordAiNumber);
                                ThrowUtils.throwIf(!save1, ErrorCode.OPERATION_ERROR, "保存数据失败");
                            }
                        } else {
                            //数据存在更新数据
                            aiAnswer[0].setMessage(JSONUtil.toJsonStr(aiMessageList));
                            boolean b = this.updateById(aiAnswer[0]);
                            ThrowUtils.throwIf(!b, ErrorCode.OPERATION_ERROR, "更新数据失败");
                            //更新用户调用AI次数记录
                            if (serviceOne != null) {
                                serviceOne.setNumber(serviceOne.getNumber() + 1);
                                boolean updateById = recordAiNumberService.updateById(serviceOne);
                                ThrowUtils.throwIf(!updateById, ErrorCode.OPERATION_ERROR, "更新数据失败");
                            }
                        }
                        //将数据加载到缓存
                        String redisKey = getRedisKey(questionId, userId);
                        redisTemplate.opsForList().rightPush(redisKey, JSONUtil.toJsonStr(new AiMessage(AiRoleEnum.USER.getValue(), userMessage)));
                        redisTemplate.opsForList().rightPush(redisKey, JSONUtil.toJsonStr(new AiMessage(AiRoleEnum.ASSISTANT.getValue(), stringBuilder.toString())));
                        //设置缓存时间为两个小时
                        redisTemplate.expire(redisKey, 120, TimeUnit.MINUTES);
                        sseEmitter.complete();
                    })
                    .subscribe();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR, "AI服务异常");
        }
        return sseEmitter;
    }

    /**
     * 根据题目Id与用户获取对话列表
     *
     * @param questionId
     * @param request
     * @return
     */
    @Override
    public List<AiMessage> getAiMessageByQuestionId(Long questionId, HttpServletRequest request) {
        //获取用户登录态
        User loginUser = userService.getLoginUser(request);
        Long userId = loginUser.getId();
        //先从缓存中加载
        String redisKey = getRedisKey(questionId, userId);
        List<String> messageList = redisTemplate.opsForList().range(redisKey, 0, -1);
        // 如果命中缓存则直接返回
        if (CollUtil.isNotEmpty(messageList)) {
            List<AiMessage> aiMessageList = new ArrayList<>();
            for (String msg : messageList) {
                aiMessageList.add(JSONUtil.toBean(msg, AiMessage.class));
            }
            log.info("命中缓存:{}",JSONUtil.toJsonStr(messageList));
            return aiMessageList;
        }else{//如果缓存为空，说明还没有对话或者缓存过期了，如果是过期了，那么将数据库的数据也删除
            //查看对话列表是否存在
            LambdaQueryWrapper<AiAnswer> queryWrapper = new LambdaQueryWrapper<>();
            queryWrapper
                    .eq(AiAnswer::getQuestionId, questionId)
                    .eq(AiAnswer::getUserId, userId);
            AiAnswer aiAnswer = this.getOne(queryWrapper);
            //如果数据不存在，直接返回空
            if (aiAnswer == null) {
                return null;
            }
            this.removeById(aiAnswer);
        }
        return null;
    }
}




